package com.xz.xzmybatisagent.socket;

import java.io.DataInputStream;
import java.io.DataOutputStream;
import java.io.IOException;
import java.net.Socket;
import java.nio.charset.StandardCharsets;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * <p>
 *
 * </p>
 *
 * @author DONG.WANG17@HAND-CHINA.COM
 * @since 2025/3/22-03-22 16:13
 */
public class SocketClient implements SocketEndpoint {
    private final SocketConfig config;
    private Socket socket;
    private DataInputStream input;
    private DataOutputStream output;
    private final ExecutorService executor = Executors.newFixedThreadPool(2);
    // 连接状态 - 已连接
    private final AtomicBoolean isConnected = new AtomicBoolean(false);
    // 是否可重试状态
    private final AtomicBoolean canRetry = new AtomicBoolean(true);
    // 重连次数
    private final AtomicInteger reconnectCount = new AtomicInteger(0);

    private final ScheduledExecutorService heartbeat = Executors.newSingleThreadScheduledExecutor();
    private final ScheduledExecutorService reconnectScheduler = Executors.newSingleThreadScheduledExecutor();
    private final BlockingQueue<String> messageQueue = new LinkedBlockingQueue<>();

    public SocketClient(SocketConfig config) {
        this.config = config;
        // 关闭钩子
        Runtime.getRuntime().addShutdownHook(new Thread(this::stop));
    }

    /**
     * 启动连接。如果连接成功，则会启动心跳线程、消息接收线程、消息发送线程。
     */
    @Override
    public void start() {
        try {
            if (isConnected.get()) return;
            executor.execute(() -> {
                while (!isConnected.get() && canRetry.get()) {
                    try {
                        // 连接
                        connect();
                        // 心跳
                        startHeartbeat();
                        // 接收监听
                        startListening();
                        // 发送监听
                        startSending();
                        // 状态通知
                        notifyStatus(ConnectionStatus.CONNECTED);
                        config.getCallback().onStart();
                    } catch (Exception e) {
                        handleError(e);
                        isConnected.set(false);
                        tryReconnect();
                        break; // 退出当前循环，等待调度任务触发
                    }
                }
            });
        } catch (Exception e) {
            handleError(e);
        }
    }

    /**
     * 停止连接。主动调用停止时不会尝试重连。
     */
    @Override
    public void stop() {
        try {
            disconnect();
            executor.shutdownNow();
            heartbeat.shutdownNow();
            reconnectScheduler.shutdownNow();
            config.getCallback().onClose();
        } catch (Exception e) {
            handleError(e);
        }
    }

    /**
     * 发送数据
     *
     * @param message 要发送的数据
     */
    @Override
    public void sendData(String message) {
        try {
            final boolean offer = messageQueue.offer(message, 1, TimeUnit.SECONDS);
            if (!offer) {
                handleError(new Exception("Message queue is full."));
            }
        } catch (Exception e) {
            handleError(e);
        }
    }

    @Override
    public boolean isRunning() {
        return isConnected.get();
    }

    /**
     * 连接服务器
     */
    private void connect() throws IOException {
        socket = new Socket(config.getHost(), config.getPort());
        input = new DataInputStream(socket.getInputStream());
        output = new DataOutputStream(socket.getOutputStream());
        isConnected.set(true);
        canRetry.set(true);
        reconnectCount.set(0);

    }

    /**
     * 启动监听线程--负责接收消息并回复（指定回复的情况下）
     */
    private void startListening() {
        executor.execute(() -> {
            while (isConnected.get()) {
                try {
                    int length = input.readInt();
                    byte[] buffer = new byte[length];
                    input.readFully(buffer);
                    String message = new String(buffer, StandardCharsets.UTF_8);

                    if (config.getHeartbeatTag().equals(message) && !config.getReceiveHeartbeatTag()) {
                        continue;
                    }
                    config.getCallback().onMessage(message);
                    final String response = config.getCallback().onMessageAndResponse(message);
                    if (response != null) {
                        // 回复消息
                        sendData(response);
                    }
                } catch (IOException e) {
                    handleError(e);
                    isConnected.set(false);
                    tryReconnect();
                    break;
                }
            }
        });
    }

    /**
     * 启动发送线程--负责发送消息
     */
    private void startSending() {
        executor.execute(() -> {
            while (isConnected.get()) {
                String message = null;
                try {
                    message = messageQueue.take();
                    byte[] data = message.getBytes(StandardCharsets.UTF_8);
                    output.writeInt(data.length);
                    output.write(data);
                    output.flush();
                } catch (Exception e) {
                    handleError(e);
                    // 发送失败，重新加入队列
                    if (message != null) {
                        messageQueue.offer(message);
                    }
                    if (canRetry.get()) {
                        tryReconnect();
                    }
                    break;
                }
            }
        });
    }

    /**
     * 启动心跳线程--负责发送心跳包
     */
    private void startHeartbeat() {
        heartbeat.scheduleAtFixedRate(() -> {
            sendData(config.getHeartbeatTag());
        }, config.getHeartbeatInitialDelay(), config.getHeartbeatPeriod(), TimeUnit.MILLISECONDS);
    }

    /**
     * 断开连接。断开操作时会尝试按照重连次数来判断是否需要断开连接。
     */
    private void disconnect() throws IOException {
        isConnected.set(false);
        if (output != null) output.close();
        if (input != null) input.close();
        if (socket != null) socket.close();
        notifyStatus(ConnectionStatus.DISCONNECTED);
    }

    /**
     * 尝试重连。重连时会进入阻塞状态
     */
    private void tryReconnect() {
        reconnectScheduler.schedule(() -> {
            canRetry.set(config.getReconnect() && reconnectCount.getAndIncrement() < config.getReconnectMaxNum());
            // 触发下一次重试检查
            if (canRetry.get()) {
                notifyStatus(ConnectionStatus.RECONNECTING);
                start(); // 重新调用启动流程
            } else {
                stop();
            }
        }, config.getReconnectPeriod(), TimeUnit.MILLISECONDS);
    }

    private void handleError(Exception e) {
        if (config.getCallback() != null) {
            config.getCallback().onError(e);
        }
    }

    private void notifyStatus(ConnectionStatus status) {
        if (config.getCallback() != null) {
            config.getCallback().onStatusChange(status);
        }
    }
}